// SPDX-License-Identifier: MPL-2.0

//! Integration layer between Rust and Zig memory management

use core::{alloc::Layout, ptr::NonNull};
use crate::mm::{
    frame::{FrameAllocOptions, GlobalFrameAllocator},
    heap::{GlobalHeapAllocator, HeapSlot, SlotInfo},
    Paddr, Vaddr, PAGE_SIZE,
};

// External Zig functions
extern "C" {
    fn zig_frame_alloc(size: usize, alignment: usize) -> *mut u8;
    fn zig_frame_dealloc(ptr: *mut u8, size: usize);
    fn zig_heap_alloc(size: usize, alignment: usize) -> *mut u8;
    fn zig_heap_dealloc(ptr: *mut u8, size: usize);
    fn zig_init_allocators(memory_start: usize, memory_size: usize);
}

// Rust functions called by Zig
#[no_mangle]
pub extern "C" fn rust_alloc_frame(size: usize, alignment: usize) -> *mut u8 {
    let layout = Layout::from_size_align(size, alignment).unwrap();
    match GlobalFrameAllocator::get().alloc(layout) {
        Some(paddr) => paddr as *mut u8,
        None => core::ptr::null_mut(),
    }
}

#[no_mangle]
pub extern "C" fn rust_dealloc_frame(ptr: *mut u8, size: usize) {
    let paddr = ptr as Paddr;
    GlobalFrameAllocator::get().dealloc(paddr, size);
}

#[no_mangle]
pub extern "C" fn rust_flush_tlb() {
    crate::arch::mm::flush_tlb_all();
}

/// Zig-based frame allocator implementation
pub struct ZigFrameAllocator {
    _private: (),
}

impl ZigFrameAllocator {
    pub fn new() -> Self {
        Self { _private: () }
    }

    pub fn init_with_memory(&self, memory_start: Paddr, memory_size: usize) {
        unsafe {
            zig_init_allocators(memory_start, memory_size);
        }
    }
}

impl GlobalFrameAllocator for ZigFrameAllocator {
    fn alloc(&self, layout: Layout) -> Option<Paddr> {
        let ptr = unsafe { zig_frame_alloc(layout.size(), layout.align()) };
        if ptr.is_null() {
            None
        } else {
            Some(ptr as Paddr)
        }
    }

    fn dealloc(&self, paddr: Paddr, size: usize) {
        unsafe {
            zig_frame_dealloc(paddr as *mut u8, size);
        }
    }
}

/// Zig-based heap allocator implementation
pub struct ZigHeapAllocator {
    _private: (),
}

impl ZigHeapAllocator {
    pub fn new() -> Self {
        Self { _private: () }
    }
}

impl GlobalHeapAllocator for ZigHeapAllocator {
    fn alloc(&self, layout: Layout) -> Result<HeapSlot, core::alloc::AllocError> {
        let ptr = unsafe { zig_heap_alloc(layout.size(), layout.align()) };
        if ptr.is_null() {
            Err(core::alloc::AllocError)
        } else {
            let slot_info = SlotInfo::from_layout(layout);
            Ok(HeapSlot::new(
                NonNull::new(ptr).unwrap(),
                slot_info.size,
                slot_info.alignment,
            ))
        }
    }

    fn dealloc(&self, slot: HeapSlot) -> Result<(), core::alloc::AllocError> {
        unsafe {
            zig_heap_dealloc(slot.as_ptr(), slot.size());
        }
        Ok(())
    }
}

/// Initialize Zig memory management with the given memory region
pub fn init_zig_memory_management(memory_start: Paddr, memory_size: usize) -> ZigFrameAllocator {
    let allocator = ZigFrameAllocator::new();
    allocator.init_with_memory(memory_start, memory_size);
    allocator
}

/// Utility function to convert between Rust and Zig page flags
pub fn rust_to_zig_page_flags(flags: crate::mm::PageFlags) -> u32 {
    let mut zig_flags = 0u32;
    
    if flags.contains(crate::mm::PageFlags::READABLE) {
        zig_flags |= 1 << 0; // Present
    }
    if flags.contains(crate::mm::PageFlags::WRITABLE) {
        zig_flags |= 1 << 1; // Writable
    }
    if flags.contains(crate::mm::PageFlags::USER_ACCESSIBLE) {
        zig_flags |= 1 << 2; // User accessible
    }
    if flags.contains(crate::mm::PageFlags::NO_CACHE) {
        zig_flags |= 1 << 4; // Cache disabled
    }
    if flags.contains(crate::mm::PageFlags::NO_EXECUTE) {
        zig_flags |= 1 << 63; // No execute (if supported)
    }
    
    zig_flags
}

/// Convert Zig page flags back to Rust page flags
pub fn zig_to_rust_page_flags(zig_flags: u32) -> crate::mm::PageFlags {
    let mut flags = crate::mm::PageFlags::empty();
    
    if zig_flags & (1 << 0) != 0 {
        flags |= crate::mm::PageFlags::READABLE;
    }
    if zig_flags & (1 << 1) != 0 {
        flags |= crate::mm::PageFlags::WRITABLE;
    }
    if zig_flags & (1 << 2) != 0 {
        flags |= crate::mm::PageFlags::USER_ACCESSIBLE;
    }
    if zig_flags & (1 << 4) != 0 {
        flags |= crate::mm::PageFlags::NO_CACHE;
    }
    if zig_flags & (1 << 63) != 0 {
        flags |= crate::mm::PageFlags::NO_EXECUTE;
    }
    
    flags
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::mm::PageFlags;

    #[test]
    fn test_page_flags_conversion() {
        let rust_flags = PageFlags::READABLE | PageFlags::WRITABLE;
        let zig_flags = rust_to_zig_page_flags(rust_flags);
        let converted_back = zig_to_rust_page_flags(zig_flags);
        
        assert_eq!(rust_flags, converted_back);
    }

    #[test]
    fn test_frame_allocator_creation() {
        let allocator = ZigFrameAllocator::new();
        // Basic creation test - actual allocation would require initialization
        assert_eq!(core::mem::size_of_val(&allocator), core::mem::size_of::<()>());
    }
}